import os
import subprocess
import json
import sys

class PackageChecker:
    def __init__(self, python_path):
        self.python_path = python_path

    def get_installed_packages(self):
        packages = []
        try:
            # 设置环境变量，强制使用 UTF-8 编码
            my_env = os.environ.copy()
            my_env["PYTHONIOENCODING"] = "utf-8"
            
            # 使用pip获取已安装软件包的详细信息
            cmd = [self.python_path, '-m', 'pip', 'list', '--format=json']
            output = subprocess.check_output(
                cmd,
                text=True,
                encoding='utf-8',
                env=my_env
            )
            
            installed_packages = json.loads(output)
            for package in installed_packages:
                size = self.get_package_size(package['name'])
                packages.append({
                    'name': package['name'],
                    'version': package.get('version', '未知'),
                    'size': size
                })
        except Exception as e:
            print(f"获取已安装软件包失败: {str(e)}")
        return packages

    def get_package_size(self, package_name):
        try:
            # 使用 pip list -v 命令获取包的位置
            cmd = [self.python_path, '-m', 'pip', 'list', '-v', '--format=json']
            my_env = os.environ.copy()
            my_env["PYTHONIOENCODING"] = "utf-8"
            output = subprocess.check_output(
                cmd,
                text=True,
                encoding='utf-8',
                env=my_env,
                stderr=subprocess.DEVNULL
            ).strip()
            
            packages_info = json.loads(output)
            location = None
            
            # 从json输出中找到对应包的位置
            for pkg in packages_info:
                if pkg['name'] == package_name:
                    location = pkg.get('location', '')
                    break
                
            if location and os.path.exists(location):
                total_size = 0
                # 尝试两种可能的包路径
                package_paths = [
                    os.path.join(location, package_name),
                    os.path.join(location, package_name.replace('-', '_'))  # 处理包名中带横线的情况
                ]
                
                for package_path in package_paths:
                    if os.path.exists(package_path):
                        for dirpath, dirnames, filenames in os.walk(package_path):
                            for f in filenames:
                                fp = os.path.join(dirpath, f)
                                if os.path.exists(fp):
                                    try:
                                        total_size += os.path.getsize(fp)
                                    except OSError:
                                        continue
                        if total_size > 0:
                            return f"{total_size / (1024*1024):.2f} MB"
                            
                # 如果没有找到包目录，尝试直接计算.dist-info或.egg-info目录的大小
                for item in os.listdir(location):
                    if item.startswith(f"{package_name.replace('-', '_')}") and (item.endswith('.dist-info') or item.endswith('.egg-info')):
                        package_path = os.path.join(location, item)
                        if os.path.exists(package_path):
                            total_size = 0
                            for dirpath, dirnames, filenames in os.walk(package_path):
                                for f in filenames:
                                    fp = os.path.join(dirpath, f)
                                    if os.path.exists(fp):
                                        try:
                                            total_size += os.path.getsize(fp)
                                        except OSError:
                                            continue
                            return f"{total_size / (1024*1024):.2f} MB"
                            
            return "未知"
        except Exception as e:
            print(f"获取包 {package_name} 大小失败: {str(e)}")
            return "未知" 